package cas.ihep.util;

import alluxio.client.file.FileSystemContext;
import alluxio.client.netty.NettyRPC;
import alluxio.network.protocol.RPCProtoMessage;
import alluxio.util.CommonUtils;
import alluxio.util.proto.ProtoMessage;
import alluxio.wire.WorkerNetAddress;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.util.concurrent.Promise;

import java.io.Closeable;
import java.io.IOException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class AlluxioRPCContext implements Closeable{

    private FileSystemContext fsContext;
    private Channel channel;
    private WorkerNetAddress workerAddress;
    private static final long default_timeout=9223372036854775807L;

    public AlluxioRPCContext(FileSystemContext ctx,WorkerNetAddress addr) throws IOException {
        fsContext=ctx;
        channel=ctx.acquireNettyChannel(workerAddress=addr);
    }

    @Override
    public void close() throws IOException {
        fsContext.releaseNettyChannel(workerAddress,channel);
    }

    private ProtoMessage internal_call(ProtoMessage request,long timeout)throws IOException{
        //Channel channel = Preconditions.checkNotNull(context.getChannel());
        final Promise<ProtoMessage> promise = channel.eventLoop().newPromise();
        channel.pipeline().addLast(new NettyRPC.RPCHandler(promise));
        channel.writeAndFlush(new RPCProtoMessage(request)).addListener(new ChannelFutureListener() {
            @Override
            public void operationComplete(ChannelFuture future) {
                if (future.cause() != null) {
                    future.channel().close();
                    promise.tryFailure(future.cause());
                }
            }
        });
        ProtoMessage message;
        try {
            message = promise.get(timeout, TimeUnit.MILLISECONDS);
        } catch (ExecutionException | TimeoutException e) {
            CommonUtils.closeChannel(channel);
            throw new IOException(e);
        } catch (InterruptedException e) {
            CommonUtils.closeChannel(channel);
            throw new RuntimeException(e);
        } finally {
            if (channel.isOpen()) {
                channel.pipeline().removeLast();
            }
        }
        if (message.isResponse()) {
            CommonUtils.unwrapResponseFrom(message.asResponse(), channel);
        }
        return message;
    }

    public ProtoMessage call(ProtoMessage request)throws IOException{
        return internal_call(request,default_timeout);
    }

    public ProtoMessage call(ProtoMessage request,long timeout)throws IOException{
        return internal_call(request, timeout);
    }
}
